# 从Milvus中检索与用户查询相似的文档分块
from config import envConfig, config
from langchain_milvus import Milvus
from langchain_ollama import OllamaEmbeddings
import logging

logger = logging.getLogger('retrieve')

def test_vectorize():
    """
    测试向量检索功能
    """
    # 配置Milvus连接
    milvusConn = {
        "uri": envConfig.MILVUS_LOCAL_URI,
        "user": envConfig.MILVUS_USER,
        "password": envConfig.MILVUS_PASSWORD
    }

    # 初始化嵌入模型
    embeddings = OllamaEmbeddings(
        model="nomic-embed-text",
        base_url=envConfig.OLLAMA_URL
    )

    # 创建检索器
    retriever = Milvus(
        collection_name=config.APP_NAME.replace(' ', '_'),
        connection_args=milvusConn,
        embedding_function=embeddings
    ).as_retriever(search_kwargs={"k": 1})  # 返回结果的数量

    # 交互式查询
    is_over = False
    while not is_over:
        query = input("请输入查询，回车查询exit退出：")
        if query == "exit":
            is_over = True
            continue
        result = retriever.get_relevant_documents(query)
        for i, ele in enumerate(result):
            print(f"{i} - Found Doc: {ele.metadata}, {ele.page_content}")

if __name__ == "__main__":
    test_vectorize()
    print("FINISH!")

